import pandas as pd
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_selection import RFECV
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
import warnings


warnings.filterwarnings("ignore")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

data = pd.read_csv('../data/train.csv')
# 取出无关列
data.drop(['EmployeeNumber', 'StandardHours', 'Over18'], axis=1, inplace=True)

# 类别列one-hot编码
categorical_cols = data.select_dtypes(include=['object']).columns.tolist()
# print(categorical_cols)
data = pd.get_dummies(data, columns=categorical_cols, drop_first=True)

x = data.drop('Attrition', axis=1)
y = data['Attrition']


# 5. 构建 Pipeline（防止数据泄露）
pipeline = Pipeline([
    ('scaler', StandardScaler()),  # 安全起见，加上（RF 不需要，但通用）
    ('rfecv', RFECV(
        estimator=RandomForestClassifier(n_estimators=50, random_state=11, n_jobs=-1),
        step=1,
        cv=StratifiedKFold(5, shuffle=True, random_state=11),
        scoring='roc_auc',
        n_jobs=-1
    ))
])

# rf = RandomForestClassifier(n_estimators=1000, random_state=11, n_jobs=-1)
#
# selector = RFECV(estimator=rf, step=1, cv=StratifiedKFold(5), scoring='roc_auc', n_jobs=-1)
# selector.fit(x, y)

pipeline.fit(x, y)
selector = pipeline.named_steps['rfecv']


print(f"最佳特征数量: {selector.n_features_}")
print(f"被选中的特征：\n{x.columns[selector.support_]}")
features_rank = pd.DataFrame({'feature': x.columns, 'rank': selector.ranking_}).sort_values('rank')
print(f"特征重要性排名：\n{features_rank}")

# 9. 可视化：CV 得分 vs 特征数量
plt.figure(figsize=(10, 6))
cv_scores = selector.cv_results_['mean_test_score']
plt.plot(range(1, len(cv_scores) + 1), cv_scores, marker='o')
plt.xlabel("选择的特征数量")
plt.ylabel("5折交叉验证的roc_auc分数")
plt.title("递归特征消除(RFECV)特征选择")
plt.grid()
plt.axvline(selector.n_features_, color='r', linestyle='--',label=f'Optimal = {selector.n_features_}')
plt.legend()
plt.tight_layout()
plt.savefig('../data/fig/RFESCV_best_features.png', dpi=300)
plt.show()